# 导入必要的库
import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# 加载预训练的GPT-2模型和分词器
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2")

def generate_text(prompt):
    # 对输入的提示进行编码
    inputs = tokenizer.encode(prompt, return_tensors="pt")

    # 生成文本
    outputs = model.generate(inputs, max_length=150, num_return_sequences=1,do_sample=False)

    # 对输出的文本进行解码
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return generated_text

# 创建Gradio界面
iface = gr.Interface(fn=generate_text, 
                     inputs="textbox", 
                     outputs="text",
                     title="GPT-2",
                     description="请先输入一段，GPT2会根据这段文本生成接下来的文本。")

# 启动应用
iface.launch()
